import numpy as np
from numpy import random as rand
import matplotlib.pyplot as plt
import matplotlib

CITY_LOC = [(41, 94), (37, 84), (54, 67), (25, 62), (7, 64),
            (2, 2), (68, 50), (71, 60), (54, 62), (83, 69)]
DNA_SIZE = len(CITY_LOC)
POP_SIZE = 10
EPOCHS = 50
CROSS_RATE = 0.7
MUTATION_RATE = 0.3


def get_distance():
    n = len(CITY_LOC)
    distance_matrix = np.zeros((n, n))
    for i, city1 in enumerate(CITY_LOC):
        for j, city2 in enumerate(CITY_LOC):
            distance_matrix[i][j] = np.sqrt((city1[0] - city2[0]) ** 2 + (city1[1] - city2[1]) ** 2)
    return distance_matrix


# 计算适应性
def get_fitness(pop, distance):
    fitness = []
    for indiv in pop:
        total_distance = 0
        n = len(indiv)
        for i in range(n):
            total_distance += distance[indiv[i], indiv[(i + 1) % n]]
        fitness.append(-total_distance)
    return fitness


# 变异
def mutation(indiv):
    if rand.rand() <= MUTATION_RATE:
        mut_idx1 = rand.randint(0, DNA_SIZE)
        mut_idx2 = rand.randint(0, DNA_SIZE)
        tmp = indiv[mut_idx1]
        indiv[mut_idx1] = indiv[mut_idx2]
        indiv[mut_idx2] = tmp


# 修复个体的基因合法性
def fix(father, mother, i1, i2):
    m_slice = mother[i1:i2 + 1]
    # 使用顺序交叉
    f_flip = np.concatenate((father[i2 + 1:], father[0:i2 + 1]))
    f_flip = np.array([x for x in f_flip if x not in m_slice])
    child = np.concatenate((f_flip[:len(father) - i2 - 1], m_slice, f_flip[len(father) - i2 - 1:]))
    return child


def select(pop, fitness):
    fitness_copy = []
    fitness_copy.extend(fitness)
    fitness_copy -= np.max(fitness_copy)
    softmax = np.exp(fitness_copy) / np.sum(np.exp(fitness_copy))
    selection = rand.choice(np.arange(len(fitness)), size=POP_SIZE, replace=True, p=softmax)
    return np.array(pop[selection])


def cross_and_mutation(pop):
    new_pop = []
    for father in pop:
        if np.random.rand() <= CROSS_RATE:
            mother = pop[rand.randint(POP_SIZE)]
            cross_idx1 = rand.randint(low=0, high=DNA_SIZE - 1)
            cross_idx2 = rand.randint(low=0, high=DNA_SIZE - 1)
            if cross_idx2 < cross_idx1:
                tmp = cross_idx1
                cross_idx1 = cross_idx2
                cross_idx2 = tmp
            child = fix(father, mother, cross_idx1, cross_idx2)
            mutation(child)
            new_pop.append(child)
        else:
            new_pop.append(father)
    return np.array(new_pop)


if __name__ == '__main__':
    matplotlib.use('TKAgg')
    np.set_printoptions(suppress=True, precision=2)
    # 距离矩阵
    distance = get_distance()
    cities = np.arange(len(CITY_LOC))
    pop = []
    # 初始种群
    for i in range(POP_SIZE):
        copy = cities.copy()
        np.random.shuffle(copy)
        pop.append(copy)
    pop = np.array(pop)
    fitness = get_fitness(pop, distance)
    for i in range(EPOCHS):
        pop = cross_and_mutation(pop)
        fitness = get_fitness(pop, distance)
        min = -np.max(fitness)
        print("epoch:{}, optimal fitness:{}".format(i, round(min, 2)))
        pop = select(pop, fitness)
    fitness = get_fitness(pop, distance)
    best_path = pop[np.argmax(fitness)]
    print(best_path)

    n = len(best_path)
    x = [0 for col in range(n + 1)]
    y = [0 for col in range(n + 1)]

    for i in range(n):
        x[i] = CITY_LOC[best_path[i]][0]
        y[i] = CITY_LOC[best_path[i]][1]
    x[n] = x[0]
    y[n] = y[0]
    plt.xlim(0, 100)
    plt.ylim(0, 100)
    plt.plot(x, y, marker='o', mec='b', label='path')
    plt.legend()
    plt.xlabel("x")
    plt.ylabel("y")
    plt.title("GA_TSP")
    plt.show()
